import matplotlib.pyplot as plt
import seaborn as sns

# Calcul des moyennes par groupe et année
trajectories = df.groupby(['year', 'income_group'])['co2_per_capita'].mean().reset_index()

import matplotlib.pyplot as plt
import seaborn as sns

# Sélection de pays à annoter
PAYS_LABELS = {
    'United States': 'USA', 'China': 'CHN', 'India': 'IND', 
    'Qatar': 'QAT', 'France': 'FRA', 'Sweden': 'SWE', 
    'Germany': 'DEU', 'Brazil': 'BRA'
}

import statsmodels.formula.api as smf
import matplotlib.pyplot as plt
import numpy as np

# Estimation avec erreurs-types robustes HC3
m_quad = smf.ols('log_co2 ~ log_gdp_c + I(log_gdp_c**2)', data=df_cross).fit(cov_type='HC3')
m_cub = smf.ols('log_co2 ~ log_gdp_c + I(log_gdp_c**2) + I(log_gdp_c**3)', data=df_cross).fit(cov_type='HC3')

# Extraction des statistiques clés
beta3 = m_cub.params['I(log_gdp_c ** 3)']
pval3 = m_cub.pvalues['I(log_gdp_c ** 3)']
aic_diff = m_quad.aic - m_cub.aic

import matplotlib.pyplot as plt
import seaborn as sns
import numpy as np

# Calcul des résidus (si modèle pas déjà estimé)
if 'm_quad' not in locals():
    m_quad = smf.ols('log_co2 ~ log_gdp_c + I(log_gdp_c**2)', data=df_cross).fit(cov_type='HC3')

# Prédictions et résidus
df_cross['pred_quad'] = m_quad.predict(df_cross)
df_cross['resid'] = df_cross['log_co2'] - df_cross['pred_quad']
df_cross['resid_abs'] = np.abs(df_cross['resid'])

# Identification des pays extrêmes
top_over = df_cross.nlargest(5, 'resid')   # Sous-estimés (polluent PLUS)
top_under = df_cross.nsmallest(5, 'resid') # Surestimés (polluent MOINS)

from sklearn.metrics import mean_squared_error
from statsmodels.tsa.arima.model import ARIMA

# Split train/test
train = df_usa[df_usa['year'] <= 2015].copy()
test = df_usa[df_usa['year'] > 2015].copy()

# Ré-estimation sur train
m_quad_train = smf.ols(
    'log_co2 ~ lag_log_co2 + log_gdp_c + I(log_gdp_c**2)', 
    data=train
).fit(cov_type='HAC', cov_kwds={'maxlags': 3})

# Prédictions
y_test = test['log_co2']
y_pred_quad = m_quad_train.predict(test)
y_pred_rw = test['lag_log_co2']  # Marche aléatoire

# ARIMA benchmark
arima_model = ARIMA(train['log_co2'], order=(1,1,1)).fit()
y_pred_arima = arima_model.forecast(steps=len(test))

# Calcul RMSE
rmse_quad = np.sqrt(mean_squared_error(y_test, y_pred_quad))
rmse_rw = np.sqrt(mean_squared_error(y_test, y_pred_rw))
rmse_arima = np.sqrt(mean_squared_error(y_test, y_pred_arima))

import matplotlib.pyplot as plt
from scipy import stats as sp_stats
from statsmodels.graphics.tsaplots import plot_acf

# ═══════════════════════════════════════════════════════════════════
# ESTIMATION DU MODÈLE (si pas déjà fait)
# ═══════════════════════════════════════════════════════════════════

if 'm_ardl_quad' not in locals():
    import statsmodels.formula.api as smf
    
    m_ardl_quad = smf.ols(
        'log_co2 ~ lag_log_co2 + log_gdp_c + I(log_gdp_c**2)', 
        data=df_usa
    ).fit(cov_type='HAC', cov_kwds={'maxlags': 3})
    
    print("⚠️ Modèle estimé dans cette cellule (devrait être fait avant)")

# Alias pour simplifier le code
m_final = m_ardl_quad

import statsmodels.formula.api as smf
import numpy as np
import pandas as pd

# Estimation du modèle ARDL quadratique (si pas déjà fait)
m_ardl_quad = smf.ols(
    'log_co2 ~ lag_log_co2 + log_gdp_c + I(log_gdp_c**2)', 
    data=df_usa
).fit(cov_type='HAC', cov_kwds={'maxlags': 3})

# Projection itérative 2021-2025
last_obs = df_usa.iloc[-1]
future_years = list(range(ANNEE_FIN + 1, ANNEE_FIN + 6))

future_df = []
curr_log_co2 = last_obs['log_co2']
curr_log_gdp = last_obs['log_gdp']

TAUX_CROISSANCE_ANNUEL = 0.02  # Hypothèse : +2%/an

for year in future_years:
    # Mise à jour du PIB
    curr_log_gdp += np.log(1 + TAUX_CROISSANCE_ANNUEL)
    curr_log_gdp_c = curr_log_gdp - GDP_MEAN_USA
    
    # Prédiction avec IC
    X_fut = pd.DataFrame({
        'lag_log_co2': [curr_log_co2],
        'log_gdp_c': [curr_log_gdp_c],
        'I(log_gdp_c ** 2)': [curr_log_gdp_c ** 2]
    })
    
    pred = m_ardl_quad.get_prediction(X_fut)
    pred_summary = pred.summary_frame(alpha=0.05)
    
    future_df.append({
        'year': year,
        'log_co2_pred': pred_summary['mean'].values[0],
        'log_co2_lower': pred_summary['mean_ci_lower'].values[0],
        'log_co2_upper': pred_summary['mean_ci_upper'].values[0],
    })
    
    curr_log_co2 = pred_summary['mean'].values[0]

df_future = pd.DataFrame(future_df)
df_future['co2_pred'] = np.exp(df_future['log_co2_pred'])
df_future['co2_lower'] = np.exp(df_future['log_co2_lower'])
df_future['co2_upper'] = np.exp(df_future['log_co2_upper'])

from statsmodels.nonparametric.smoothers_lowess import lowess
import matplotlib.pyplot as plt
import numpy as np

# Modèles paramétriques (si pas déjà estimés)
if 'm_quad' not in locals():
    m_quad = smf.ols('log_co2 ~ log_gdp_c + I(log_gdp_c**2)', data=df_cross).fit(cov_type='HC3')
    m_cub = smf.ols('log_co2 ~ log_gdp_c + I(log_gdp_c**2) + I(log_gdp_c**3)', data=df_cross).fit(cov_type='HC3')

# LOWESS (non-paramétrique)
lowess_result = lowess(df_cross['log_co2'], df_cross['log_gdp_c'], frac=0.3)